from sqlalchemy import create_engine, Column, Integer, String, Float, DateTime
from sqlalchemy.orm import declarative_base, sessionmaker
from datetime import datetime
import sys
import os

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 创建简化的模型以用于数据初始化
Base = declarative_base()

class DroneType(Base):
    __tablename__ = "drone_types"
    
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String(100), unique=True, index=True, nullable=False)
    description = Column(String(255))
    max_speed = Column(Float)
    max_payload = Column(Float)
    max_flight_time = Column(Integer)  # 分钟
    created_at = Column(DateTime, default=datetime.now)
    updated_at = Column(DateTime, default=datetime.now, onupdate=datetime.now)

# 数据库连接
DATABASE_URL = "sqlite:///./pilot_manage.db"
engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

# 从DroneTypes.vue中提取的无人机型号数据
drone_types_data = [
    {
        "id": 1,
        "name": "DJI Mavic 3",
        "brand": "DJI",
        "max_flight_time": 46,
        "max_flight_distance": 18,
        "max_speed": 72,
        "weight": 895,
        "created_at": "2024-01-01 00:00:00",
        "description": "高端专业航拍无人机"
    },
    {
        "id": 2,
        "name": "DJI Mini 3 Pro",
        "brand": "DJI",
        "max_flight_time": 34,
        "max_flight_distance": 10,
        "max_speed": 65,
        "weight": 249,
        "created_at": "2024-01-02 10:00:00",
        "description": "便携航拍无人机，重量低于250g"
    },
    {
        "id": 3,
        "name": "DJI Air 3",
        "brand": "DJI",
        "max_flight_time": 46,
        "max_flight_distance": 18,
        "max_speed": 72,
        "weight": 720,
        "created_at": "2024-01-03 14:00:00",
        "description": "中高端航拍无人机"
    },
    {
        "id": 4,
        "name": "DJI Phantom 4",
        "brand": "DJI",
        "max_flight_time": 28,
        "max_flight_distance": 7,
        "max_speed": 72,
        "weight": 1380,
        "created_at": "2024-01-04 09:00:00",
        "description": "经典专业航拍无人机"
    },
    {
        "id": 5,
        "name": "Autel Evo II",
        "brand": "Autel",
        "max_flight_time": 40,
        "max_flight_distance": 9,
        "max_speed": 72,
        "weight": 1100,
        "created_at": "2024-01-05 16:00:00",
        "description": "专业级航拍无人机"
    },
    {
        "id": 6,
        "name": "Skydio 2+",
        "brand": "Skydio",
        "max_flight_time": 23,
        "max_flight_distance": 6,
        "max_speed": 36,
        "weight": 775,
        "created_at": "2024-01-06 11:00:00",
        "description": "智能避障无人机"
    }
]

def init_drone_types():
    db = SessionLocal()
    try:
        # 检查是否已存在无人机型号数据
        existing_types_count = db.query(DroneType).count()
        if existing_types_count > 0:
            print(f"数据库中已存在 {existing_types_count} 种无人机型号数据，跳过初始化")
            return
        
        # 初始化无人机型号数据
        for type_data in drone_types_data:
            # 对于数据库中不存在的字段（brand、max_flight_distance、weight），我们可以在description中包含这些信息
            enhanced_description = f"{type_data['description']}。品牌：{type_data['brand']}，重量：{type_data['weight']}g"
            
            drone_type = DroneType(
                id=type_data["id"],
                name=type_data["name"],
                description=enhanced_description,
                max_speed=type_data["max_speed"],
                max_payload=type_data["weight"] / 1000,  # 将重量转换为大致的载重能力（kg）
                max_flight_time=type_data["max_flight_time"],
                created_at=datetime.strptime(type_data["created_at"], "%Y-%m-%d %H:%M:%S")
            )
            db.add(drone_type)
            print(f"添加无人机型号: {type_data['name']}")
        
        db.commit()
        print("成功初始化 6 种无人机型号数据到数据库")
    except Exception as e:
        db.rollback()
        print(f"初始化无人机型号数据时出错: {e}")
    finally:
        db.close()

if __name__ == "__main__":
    # 创建表（如果不存在）
    Base.metadata.create_all(bind=engine)
    
    # 初始化数据
    init_drone_types()